#include "vmm.h"
#include "printf.h"
#include "pmm.h"  // 确保引入pmm的页分配/释放函数声明
#include <stddef.h>

// 全局内核页表（与vmm.h中声明对应）
pagetable_t kernel_pagetable = NULL;

// 创建新页表（分配一页物理内存并清零）
pagetable_t create_pagetable(void) {
    void* page = alloc_page();  // 依赖pmm提供的单页分配函数
    if (!page) {
        printf("创建页表：分配物理页失败\n");
        return NULL;
    }
    
    // 初始化页表所有项为0（无效）
    pagetable_t pt = (pagetable_t)page;
    for (int i = 0; i < 512; i++) {  // Sv39每页表含512个PTE
        pt[i] = 0;
    }
    printf("创建页表：新页表位于物理地址 PA %p\n", (void*)pt);
    return pt;
}

// 查找虚拟地址对应的页表项，不存在则创建中间页表
pte_t* walk_create(pagetable_t pt, uint64_t va) {
    if (!pt) {
        printf("页表查找（创建）：无效的页表\n");
        return NULL;
    }

    // 遍历三级页表（L2→L1→L0，Sv39模式）
    for (int level = PG_LEVELS - 1; level > 0; level--) {
        uint64_t idx = VPN_MASK(va, level);  // 提取当前级VPN索引（9位）
        pte_t pte = pt[idx];
        
        // 页表项无效：创建新的中间页表
        if (!(pte & PTE_V)) {
            pagetable_t new_pt = create_pagetable();
            if (!new_pt) {
                printf("页表查找（创建）：在级别 %d 创建中间页表失败\n", level);
                return NULL;
            }
            
            // 中间页表项权限：有效（V）+ 可读（R）+ 可写（W）（需修改下一级页表）
            pt[idx] = PA2PTE((uint64_t)new_pt) | PTE_V | PTE_R | PTE_W;
        }
        
        // 进入下一级页表（物理地址转换）
        pt = (pagetable_t)PTE2PA(pt[idx]);
    }

    // 返回最后一级（L0）页表中对应的PTE指针
    return &pt[VPN_MASK(va, 0)];
}

// 查找虚拟地址对应的页表项（仅查找，不创建）
pte_t* walk_lookup(pagetable_t pt, uint64_t va) {
    if (!pt) {
        printf("页表查找（仅查询）：无效的页表\n");
        return NULL;
    }

    // 遍历三级页表
    for (int level = PG_LEVELS - 1; level > 0; level--) {
        uint64_t idx = VPN_MASK(va, level);
        pte_t pte = pt[idx];
        
        // 页表项无效，查找失败
        if (!(pte & PTE_V)) {
            printf("页表查找（仅查询）：级别 %d 的页表项无效\n", level);
            return NULL;
        }
        
        // 进入下一级页表
        pt = (pagetable_t)PTE2PA(pt[idx]);
    }

    // 返回L0级PTE指针
    return &pt[VPN_MASK(va, 0)];
}

// 建立单个虚拟页到物理页的映射
int map_page(pagetable_t pt, uint64_t va, uint64_t pa, int perm) {
    // 检查地址是否按页对齐（PGSIZE=4096）
    if ((va % PGSIZE) != 0 || (pa % PGSIZE) != 0) {
        printf("映射页面：虚拟地址 va=%p 或物理地址 pa=%p 未按页对齐\n", (void*)va, (void*)pa);
        return -1;  // 地址未对齐
    }

    // 查找或创建页表项
    pte_t* pte = walk_create(pt, va);
    if (!pte) {
        printf("映射页面：为虚拟地址 va=%p 创建页表项失败\n", (void*)va);
        return -2;  // 内存分配失败
    }

    // 检查是否已存在有效映射
    if (*pte & PTE_V) {
        printf("映射页面：虚拟地址 va=%p 已存在映射（冲突）\n", (void*)va);
        return -3;  // 映射冲突
    }

    // 设置页表项：物理地址+权限+有效位
    *pte = PA2PTE(pa) | perm | PTE_V;
    printf("映射页面：虚拟地址 va=%p -> 物理地址 pa=%p （权限：R=%d W=%d X=%d U=%d）\n",
           (void*)va, (void*)pa,
           (perm & PTE_R) ? 1 : 0,
           (perm & PTE_W) ? 1 : 0,
           (perm & PTE_X) ? 1 : 0,
           (perm & PTE_U) ? 1 : 0);
    return 0;  // 映射成功
}

// 映射连续虚拟地址区域到连续物理地址区域
void map_region(pagetable_t pt, uint64_t va, uint64_t pa, uint64_t size, int perm) {
    // 检查地址和大小是否页对齐
    if ((va % PGSIZE != 0) || (pa % PGSIZE != 0) || (size % PGSIZE != 0)) {
        panic("映射区域：地址或大小未按页对齐");
    }
    if (size == 0) {
        printf("映射区域：大小为0，跳过\n");
        return;
    }

    // 逐页映射
    uint64_t pages = size / PGSIZE;
    for (uint64_t i = 0; i < pages; i++) {
        int ret = map_page(pt, va + i * PGSIZE, pa + i * PGSIZE, perm);
        if (ret != 0) {
            panic("映射区域：第 %ld 页映射失败（返回值=%d）", i, ret);
        }
    }
    printf("映射区域：虚拟地址范围 va=%p~%p -> 物理地址范围 pa=%p~%p （大小=%ldKB，权限：R=%d W=%d X=%d）\n",
           (void*)va, (void*)(va + size),
           (void*)pa, (void*)(pa + size),
           size / 1024,
           (perm & PTE_R) ? 1 : 0,
           (perm & PTE_W) ? 1 : 0,
           (perm & PTE_X) ? 1 : 0);
}

// 递归销毁页表
static void destroy_pagetable_recursive(pagetable_t pt, int level) {
    if (!pt) return;

    // 遍历当前级页表的所有PTE
    for (int i = 0; i < 512; i++) {
        pte_t pte = pt[i];
        if (pte & PTE_V) {
            uint64_t pa = PTE2PA(pte);
            // 中间页表（非最后一级且有读写权限）需递归销毁
            if (level < PG_LEVELS - 1 && (pte & (PTE_R | PTE_W))) {
                destroy_pagetable_recursive((pagetable_t)pa, level + 1);
            }
        }
    }

    // 释放当前页表占用的物理页
    free_page(pt);
    printf("递归销毁页表：释放级别 %d 的页表，位于物理地址 PA %p\n", level, pt);
}

// 销毁页表入口函数
void destroy_pagetable(pagetable_t pt) {
    if (pt) {
        destroy_pagetable_recursive(pt, 0);
    } else {
        printf("销毁页表：无效的页表\n");
    }
}

// 递归打印页表内容
static void dump_pagetable_recursive(pagetable_t pt, int level, uint64_t va) {
    if (!pt) return;

    for (int i = 0; i < 512; i++) {
        pte_t pte = pt[i];
        if (pte & PTE_V) {
            uint64_t current_va = va | ((uint64_t)i << VPN_SHIFT(level));
            uint64_t pa = PTE2PA(pte);

            // 缩进显示页表级别（L2→L1→L0递进）
            for (int j = 0; j < level; j++) printf("  ");
            printf("L%d：VPN=%03lx -> 物理地址 PA=%p | ", level, i, (void*)pa);
            printf("V=%d R=%d W=%d X=%d U=%d A=%d D=%d\n",
                   (pte & PTE_V) ? 1 : 0,
                   (pte & PTE_R) ? 1 : 0,
                   (pte & PTE_W) ? 1 : 0,
                   (pte & PTE_X) ? 1 : 0,
                   (pte & PTE_U) ? 1 : 0,
                   (pte & PTE_A) ? 1 : 0,
                   (pte & PTE_D) ? 1 : 0);

            // 递归打印下一级页表（中间页表）
            if (level < PG_LEVELS - 1 && (pte & (PTE_R | PTE_W))) {
                dump_pagetable_recursive((pagetable_t)pa, level + 1, current_va);
            }
        }
    }
}

// 打印页表内容入口函数
void dump_pagetable(pagetable_t pt, int level) {
    if (!pt) {
        printf("打印页表：无效的页表\n");
        return;
    }
    printf("===== 页表内容打印（开始） =====");
    if (level >= 0) printf("（级别 %d）", level);
    printf("\n");
    dump_pagetable_recursive(pt, level, 0);
    printf("===== 页表内容打印（结束） =====\n");
}

// 激活当前核心的虚拟内存（写入SATP寄存器并刷新TLB）
static void kvm_inithart(void) {
    if (!kernel_pagetable) {
        panic("内核虚拟内存初始化（hart）：内核页表为空");
    }

    // 构造SATP寄存器值（Sv39模式 + 根页表物理页号）
    uint64_t satp_val = MAKE_SATP(kernel_pagetable);
    printf("内核虚拟内存初始化（hart）：设置 SATP 寄存器为 0x%lx\n", satp_val);

    // 写入SATP寄存器（RISC-V控制寄存器操作）
    asm volatile("csrw satp, %0" : : "r"(satp_val));
    // 刷新TLB（确保旧映射失效）
    asm volatile("sfence.vma");
    printf("内核虚拟内存初始化（hart）：虚拟内存已启用（SATP 已更新）\n");
}

// 初始化内核虚拟内存（创建页表并完成所有映射）
void kvm_init(void) {
    // 打印虚拟内存初始化标题（与其他模块风格统一）
    printf("=== 虚拟内存初始化 ===\n");

    // 1. 创建内核根页表（依赖PMM已初始化）
    kernel_pagetable = create_pagetable();
    if (!kernel_pagetable) {
        panic("kvm_init：创建内核根页表失败");
    }

    // 2. 映射内核代码段（.text段，R+X权限）
    uint64_t text_pa_start = (uint64_t)_text;
    uint64_t text_pa_end = (uint64_t)_etext;
    uint64_t text_va_start = KERNBASE + (text_pa_start - PHYSBASE);  // 虚拟地址 = 偏移 + 物理地址
    uint64_t text_size = text_pa_end - text_pa_start;
    map_region(kernel_pagetable, text_va_start, text_pa_start, text_size, PTE_R | PTE_X);

    // 3. 映射内核数据段（.data/.bss等，R+W权限）
    uint64_t data_pa_start = text_pa_end;
    uint64_t data_pa_end = (uint64_t)_end;
    uint64_t data_va_start = KERNBASE + (data_pa_start - PHYSBASE);
    uint64_t data_size = data_pa_end - data_pa_start;
    map_region(kernel_pagetable, data_va_start, data_pa_start, data_size, PTE_R | PTE_W);

    // 4. 映射内核页表自身（启用虚拟内存后需通过虚拟地址访问页表）
    uint64_t pt_pa = (uint64_t)kernel_pagetable;
    uint64_t pt_va = KERNBASE + (pt_pa - PHYSBASE);  // 页表虚拟地址 = 偏移 + 物理地址
    map_region(kernel_pagetable, pt_va, pt_pa, PGSIZE, PTE_R | PTE_W);

    // 5. 映射UART设备（虚拟地址=物理地址，R+W权限）
    map_region(kernel_pagetable, UART0, UART0, PGSIZE, PTE_R | PTE_W);

    // 6. 激活虚拟内存（单核心场景）
    kvm_inithart();
}
